import copy
import os
import time
import torch
import torch.nn as nn
import pandas as pd
import dataset
from model import FCN8s
import utils

def train_process(model, dataTrain, dataVal, epoch, load_from_checkpoint=False, checkpoint_path='../checkpoints/best_model.pth'):
    checkpoint_dir = os.path.dirname(checkpoint_path)  # 取目录
    if not os.path.exists(checkpoint_dir):
        os.makedirs(checkpoint_dir)
        print(f"[INFO] 目录 {checkpoint_dir} 不存在，已创建。")

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    if load_from_checkpoint and os.path.exists(checkpoint_path):
        print(f"-> 加载 模型权重文件: {checkpoint_path} 继续训练...")
        checkpoint = torch.load(checkpoint_path, map_location=device)

        model.load_state_dict(checkpoint['model_state_dict'])
        best_model_wts = copy.deepcopy(model.state_dict())

        best_val_loss = checkpoint['best_val_loss']

        print(f"当前权loss值:{best_val_loss}")

    else:
        print("-> 从头开始训练 ")
        best_model_wts = copy.deepcopy(model.state_dict())
        best_val_loss = float('inf')



    lossFunc = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    accTrainAll, accValAll = [], []
    lossTrainAll, lossValAll = [], []

    since = time.time()

    for i in range(epoch):
        print(f"=== 第 {i+1}/{epoch} 轮训练 ===")
        # 训练
        model.train()
        running_loss_train = 0
        running_corrects_train = 0
        total_train = 0

        for imgs, masks in dataTrain:
            imgs = imgs.to(device)
            masks = masks.to(device)

            optimizer.zero_grad()
            outputs = model(imgs)
            loss = lossFunc(outputs, masks)
            loss.backward()
            optimizer.step()

            running_loss_train += loss.item() * imgs.size(0)
            preds = torch.argmax(outputs, dim=1)
            running_corrects_train += torch.sum(preds == masks).item()
            total_train += masks.numel()  # 每个像素算一条

        epoch_loss_train = running_loss_train / len(dataTrain.dataset)
        epoch_acc_train = running_corrects_train / total_train

        # 验证
        model.eval()
        running_loss_val = 0
        running_corrects_val = 0
        total_val = 0

        with torch.no_grad():
            for imgs, masks in dataVal:
                imgs = imgs.to(device)
                masks = masks.to(device)

                outputs = model(imgs)
                loss = lossFunc(outputs, masks)

                running_loss_val += loss.item() * imgs.size(0)
                preds = torch.argmax(outputs, dim=1)
                running_corrects_val += torch.sum(preds == masks).item()
                total_val += masks.numel()

        epoch_loss_val = running_loss_val / len(dataVal.dataset)
        epoch_acc_val = running_corrects_val / total_val

        accTrainAll.append(epoch_acc_train)
        accValAll.append(epoch_acc_val)
        lossTrainAll.append(epoch_loss_train)
        lossValAll.append(epoch_loss_val)

        # 保存最佳模型
        if epoch_loss_val < best_val_loss:
            best_val_loss = epoch_loss_val
            best_model_wts = copy.deepcopy(model.state_dict())
            torch.save({'model_state_dict': best_model_wts, 'best_val_loss': best_val_loss}, checkpoint_path)
            print(f"-> 产生新的最好权重：  {checkpoint_path}")



        time_elapsed = time.time() - since
        print(f"训练耗时: {time_elapsed//60:.0f}m{time_elapsed%60:.0f}s")
        print(f"训练集准确率: {epoch_acc_train:.4f}, 验证集准确率: {epoch_acc_val:.4f}")
        print(f"训练集Loss: {epoch_loss_train:.4f}, 验证集Loss: {epoch_loss_val:.4f}")
        print("-"*30)

    # 画图
    draw_data = pd.DataFrame({
        "epoch": range(1, epoch+1),
        "accTrain": accTrainAll,
        "accVal": accValAll,
        "lossTrain": lossTrainAll,
        "lossVal": lossValAll
    })
    return draw_data

if __name__ == '__main__':
    # 获取数据集
    dataset_train_path = "../Portrait-dataset-2000/dataset/training"
    dataset_val_path = "../Portrait-dataset-2000/dataset/testing"
    epoch = 10
    batch_size = 32
    num_workers = 4
    load_from_checkpoint =  False
    checkpoint_path = '../checkpoints/best_model.pth'

    dataset = dataset.DatasetProcess(dataset_train_path, dataset_val_path, batch_size, num_workers)
    dataTrain, dataVal = dataset.get_data()
    model = FCN8s(num_classes=2) # 背景、人
    draw_data = train_process(model=model, dataTrain=dataTrain, dataVal=dataVal, epoch=epoch, load_from_checkpoint=load_from_checkpoint, checkpoint_path=checkpoint_path)

    utils.draw(draw_data)

